import argparse
import logging
import sys
import threading
import time
from concurrent import futures

import cv2
import grpc
import numpy as np
import zmq

import updater
from lite import FaceRecognitionLite

sys.path.append(f'{sys.path[0]}/pb')

from pb import to_pb2, task_pb2, to_pb2_grpc, aux_task_pb2

conf = None
lock = threading.Lock()

lite: FaceRecognitionLite
socket_push: zmq.Socket
socket_sub: zmq.Socket


class MySendMsgTOAIServicer(to_pb2_grpc.SendMsgTOAIServicer):
    def SendMsgTOAI(self, req, ctx):
        # sync start
        lock.acquire()

        start = round(time.time() * 1e3)
        logging.debug('[FR Lite] Receive request')

        basic = req.basic
        latency = req.latency

        res = aux_task_pb2.Res()

        if basic.taskType != aux_task_pb2.Basic.EP:

            # obtain image
            img_bytes = basic.snapshot
            img = cv2.imdecode(np.asarray(bytearray(img_bytes)), cv2.IMREAD_COLOR)
            logging.debug(f'Got img of shape ({img.shape})')

            # execute main FR
            persons, gender, age, traffic, img = lite.new_frame(img).lite()
            logging.debug(f'[FR Lite] Read lite: {persons}, {gender}, {age}, {traffic}')

            # classify personType
            tp_cnt = {
                'Student': 0,
                'Teacher': 0,
                'Visitor': 0,
                'Stranger': 0,
                'Blacklist': 0
            }
            for person in persons:
                # set type
                tp_cnt[person['type']] += 1
                # add person to proto
                p = res.persons.add()
                p.name = person['name']
                p.type = getattr(aux_task_pb2.Res.Person, person['type'])
                for val in person['rectangle']:
                    p.rectangle.append(val)

            # set alert
            res.alert.append(tp_cnt['Blacklist'])
            res.alert.append(tp_cnt['Stranger'])
            # set flow
            res.flow.append(tp_cnt['Student'])
            res.flow.append(tp_cnt['Teacher'])
            res.flow.append(tp_cnt['Visitor'])
            # set gender
            res.gender.append(gender['Male'])
            res.gender.append(gender['Female'])

        else:
            logging.debug('[FR Lite] Empty Request')

        # set timestamps
        end = round(time.time() * 1e3)
        latency.timestamps.append(start)
        latency.timestamps.append(end)

        # send result
        task = task_pb2.Task(basic=basic, latency=latency, res=res)
        socket_push.send(task.SerializeToString())
        logging.debug('[FR Lite] Sending result')
        # sync end
        lock.release()

        return to_pb2.BackTO(res=res)


def renew_db(ip, port):
    is_data = False
    while True:
        socket_sub.recv()
        if is_data:
            updater.update_db(ip, port)
            lite.renew_db()
            logging.info("[FR DB] new DB loaded")
        is_data = ~is_data


def init(args):
    global conf, lite, socket_push, socket_sub
    # basic settings
    conf = args
    logging.basicConfig(
        level=conf.log_level,
        format='[%(asctime)s] <%(levelname)s> %(filename)s:%(lineno)d: %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S')
    logging.info(f'[FR Main] ds_addr: {conf.ds_addr}, ai_port: {conf.ai_port}')
    logging.info('[FR Main] Checking files...')
    # update FR db and model
    ds_ip = str(conf.ds_addr).split(":")[0]
    updater.check_file(ds_ip, conf.bed_port)
    # init FR processor
    lite = FaceRecognitionLite(None, conf.detect_mode, conf.do_option, conf.debug_mode)
    # init zmq (pusher, ai->ds)
    logging.info('[FR Main] Init pusher...')
    context = zmq.Context()
    socket_push = context.socket(zmq.PUSH)
    socket_push.connect("tcp://{}".format(conf.ds_addr))
    logging.info('[FR Main] Zmq pusher started')
    # init zmq (subscriber, db->lite)
    logging.info('[FR Main] Init subscriber...')
    socket_sub = context.socket(zmq.SUB)
    socket_sub.connect("tcp://{}:{}".format(ds_ip, conf.db_port))
    socket_sub.setsockopt(zmq.SUBSCRIBE, bytes(conf.db_topic, 'utf-8'))
    threading.Thread(target=renew_db, args=(ds_ip, conf.bed_port)).start()
    logging.info('[FR Main] Zmq subscriber started')
    # init grpc (to->ai)
    logging.info('[FR Main] Booting server...')
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    to_pb2_grpc.add_SendMsgTOAIServicer_to_server(
        MySendMsgTOAIServicer(), server)
    server.add_insecure_port("0.0.0.0:{}".format(conf.ai_port))
    server.start()
    logging.info('[FR Main] gRPC server started')

    try:
        server.wait_for_termination()
    except KeyboardInterrupt:
        logging.info('[FR Main] Server ended')


def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('-ds', '--ds_addr', type=str, default='localhost:5800',
                        help='ZeroMQ puller address of Dashboard.')
    parser.add_argument('-ai', '--ai_port', type=int, default=6585,
                        help='Grpc server port of AI.')
    parser.add_argument('--bed_port', type=int, default=5000,
                        help='ZeroMQ publisher port of PictureBed.')
    parser.add_argument('--db_port', type=int, default=9000,
                        help='ZeroMQ publisher port of FRLite-DB.')
    parser.add_argument('--db_topic', type=str, default='db',
                        help='ZeroMQ publisher topic of FRLite-DB.')
    parser.add_argument('--log_level', type=str, choices=['DEBUG', 'INFO', 'WARN', 'ERROR', 'FATAL'],
                        help='ZeroMQ publisher port of PictureBed.', default='DEBUG')
    parser.add_argument('--detect_mode', type=str, choices=['CASCADE', 'MODEL'],
                        help='Face detection mode.', default='MODEL')
    parser.add_argument('--do_option', action="store_true",
                        help='whether to calculate optional attributes (traffic) or not.')
    parser.add_argument('--debug_mode', action="store_true")

    return parser.parse_args(argv)


if __name__ == '__main__':
    init(parse_arguments(sys.argv[1:]))
